"""
Copyright © 2024 Walkline Wang (https://walkline.wang)
Gitee: https://gitee.com/walkline/micropython-batch-ble-config-device
"""
import time
import network
import bluetooth
import smartconfig
from machine import reset
from blelib.ble import BLETools, IRQ
from drivers.button import Button

import sys
sys.path.append('/blelib')

import esp
esp.osdebug(None) # 注释此行可显示详细调试信息


MODE_SERVER = 0
MODE_CLIENT = 1

TARGET_DEVICE = 'ble-config'

SSID_PREFIX		= 'ssid_'
PASSWORD_PREFIX	= 'pswd_'

STA_CONFIG_FILENAME    = 'sta_config.py'
STA_CONFIG_IMPORT_NAME = STA_CONFIG_FILENAME.split('.')[0]

NORDIC_UART_UUID = bluetooth.UUID('6E400001-B5A3-F393-E0A9-E50E24DCCA9E')
NORDIC_RX_UUID   = bluetooth.UUID('6E400002-B5A3-F393-E0A9-E50E24DCCA9E')
NORDIC_TX_UUID   = bluetooth.UUID('6E400003-B5A3-F393-E0A9-E50E24DCCA9E')

def printf(msg, *args, **kwargs):
	print(f'\033[1;37m[INFO]\033[0m {msg}', *args, **kwargs)


class FileUtils(object):
	@staticmethod
	def output_sta_config_file(ssid, password):
		with open(STA_CONFIG_FILENAME, 'w') as output:
			output.write(
f'''# automatic generated file
ssid = '{ssid}'
password = '{password}'
''')

	@staticmethod
	def delete_sta_config_file():
		import os
		try:
			os.remove(STA_CONFIG_FILENAME)
		except:
			pass


class Utilities(object):
	SC_TIMEOUT = 120_000 # ms

	@staticmethod
	def choose_an_option(title: str, options: list | tuple) -> (int | None):
		print(f'\n{title}:')

		for index, option in enumerate(options, start=1):
			if index == 1:
				print(f'\x1b[32m  [{index}] {option}\033[0m')
			else:
				print(f'  [{index}] {option}')

		selected = None

		while True:
			try:
				selected = input('Choose an option: ')

				if selected == '':
					return 0

				selected = int(selected)

				assert type(selected) is int and 0 < selected <= len(options)

				return selected - 1
			except KeyboardInterrupt:
				return
			except:
				pass

	@staticmethod
	def start_smartconfig() -> bool:
		printf('Starting SmartConfig...')

		result = False

		sta = network.WLAN(network.STA_IF)
		_ = sta.active(True)

		smartconfig.type(smartconfig.TYPE_ESPTOUCH_AIRKISS)
		smartconfig.start()

		start_ms = time.ticks_ms()

		while not smartconfig.done():
			if time.ticks_ms() - start_ms > Utilities.SC_TIMEOUT:
				printf('SmartConfig Timedout')
				smartconfig.stop()
				return

			time.sleep_ms(100)

		print('SmartConfig done, ', end='')

		if sta.status() == network.STAT_GOT_IP:
			print('and connected to ap')
			printf(f'SmartConfig info: {smartconfig.info()}')
			printf(f'  - ssid: "{smartconfig.ssid()}"')
			printf(f'  - password: "{smartconfig.password()}"')
			printf(f'  - bssid: {smartconfig.bssid()}')
			printf(f'  - type: {smartconfig.type()}')

			if smartconfig.rvd_data():
				printf(f'  - rvd_data: {smartconfig.rvd_data()}') # EspTouch V2 custom data

			result = True
		else:
			# maybe wrong password or other situations
			print('but failed connect to ap')

		smartconfig.stop()

		return result

	@staticmethod
	def is_esp32c3() -> bool:
		from machine import Pin
		from os import uname

		if uname()[0] == 'esp32':
			try:
				Pin(22, Pin.OUT)
				return False
			except ValueError:
				return True

	@staticmethod
	def hard_reset():
		network.WLAN().active(False)
		reset()


class Client(object):
	def __init__(self, data):
		self.addr_type  = data[0]
		self.addr       = bytes(data[1])
		self.adv_type   = data[2]
		self.__rssi     = data[3]
		self.__adv_data = bytes(data[4])

		self.name          = BLETools.decode_name(self.__adv_data) or BLETools.decode_mac(self.addr)
		self.conn_handle   = None
		self.found         = False
		self.start_handle  = 0
		self.end_handle    = 0

		self.handle_rx = None


class Scanner(object):
	def __init__(self,
			device_name: str = 'ble-scanner',
			target_device: str = None,
			found_rx_cb: function = None):
		self.__ble = bluetooth.BLE()
		self.__client: Client = None

		self.__target_device = target_device
		self.__found_rx_cb   = found_rx_cb

		self.__write = self.__ble.gattc_write

		self.__ble.config(gap_name=device_name)
		self.__ble.irq(self.__irq_callback)

		self.__ble.active(False)
		printf('Activating BLE...')
		self.__ble.active(True)
		printf(f'BLE Activated [{BLETools.decode_mac(self.__ble.config('mac')[1])}]')

		self.__ble.config(mtu=256)

	def __irq_callback(self, event, data):
		if event == IRQ.PERIPHERAL_CONNECT:
			conn_handle, addr_type, addr = data

			printf(f'[{self.__client.name}] Connected')
			self.__client.conn_handle = conn_handle
			self.__ble.gattc_discover_services(conn_handle)
		elif event == IRQ.PERIPHERAL_DISCONNECT:
			if self.__client:
				printf(f'[{self.__client.name}] Disconnected')
				self.__client.conn_handle = None
		elif event == IRQ.SCAN_RESULT:
			addr_type, addr, adv_type, rssi, adv_data = data

			if self.__check(adv_data):
				self.__client = Client(data)
				self.__ble.gap_scan(None)
		elif event == IRQ.SCAN_DONE:
			printf(f'Scan completed, BLE config client{"" if self.__client else " not"} found')

			if self.__client:
				printf(f'Connecting to [{self.__client.name}]')
				self.__ble.gap_connect(self.__client.addr_type, self.__client.addr)
		elif event == IRQ.GATTC_SERVICE_RESULT:
			conn_handle, start_handle, end_handle, uuid = data

			if uuid == NORDIC_UART_UUID:
				self.__client.found = True
				self.__client.start_handle  = start_handle
				self.__client.end_handle    = end_handle
		elif event == IRQ.GATTC_SERVICE_DONE:
			conn_handle, status = data

			if self.__client.found:
				self.__ble.gattc_discover_characteristics(conn_handle, self.__client.start_handle, self.__client.end_handle)
		elif event == IRQ.GATTC_CHARACTERISTIC_RESULT:
			conn_handle, end_handle, value_handle, properties, uuid = data

			if uuid == NORDIC_RX_UUID:
				self.__client.handle_rx = value_handle
		elif event == IRQ.GATTC_CHARACTERISTIC_DONE:
			if self.__client.handle_rx:
				if self.__found_rx_cb is not None:
					self.__found_rx_cb()

	def __check(self, adv_data):
		return BLETools.decode_name(adv_data) == self.__target_device

	def scan(self):
		self.__client = None
		printf(f'Scaning...')
		self.__ble.gap_scan(0, 50000, 50000, True)

	def write(self, data: bytes):
		if self.__client and self.__client.conn_handle and self.__client.handle_rx:
			self.__write(self.__client.conn_handle, self.__client.handle_rx, data)

	def disconnect(self):
		if self.__client:
			self.__ble.gap_disconnect(self.__client.conn_handle)


def run_bleconfig_server():
	try:
		sta_config = __import__(STA_CONFIG_IMPORT_NAME)
	except ImportError:
		if Utilities.start_smartconfig():
			FileUtils.output_sta_config_file(smartconfig.ssid(), smartconfig.password())
		else:
			Utilities.hard_reset()

	def found_rx_cb():
		printf('Sending data')
		scanner.write(f'{SSID_PREFIX}{ssid}'.encode())
		scanner.write(f'{PASSWORD_PREFIX}{password}'.encode())
		printf('BLE config completed')
		scanner.disconnect()
		scanner.scan()

	sta_config = __import__(STA_CONFIG_IMPORT_NAME)
	ssid       = sta_config.ssid
	password   = sta_config.password

	scanner = Scanner(target_device=TARGET_DEVICE, found_rx_cb=found_rx_cb)
	scanner.scan()

def run_bleconfig_client():
	from blelib.devices.uart.bleconfig import BLEConfig

	def rx_received_cb(data: bytes):
		printf(f'Received Data: {data}')

	bleconfig = BLEConfig(device_name=TARGET_DEVICE, rx_received_cb=rx_received_cb)

	while not bleconfig.success():
		time.sleep_ms(100)

	printf(f'ssid: {bleconfig.ssid}, password: {bleconfig.password}')

def run_test():
	def button_press_callback(pin:int):
		printf(f'Button {pin} pressed')

		FileUtils.delete_sta_config_file()
		Utilities.hard_reset()

	button = Button(
		pin=9 if Utilities.is_esp32c3() else 0,
		press_cb=button_press_callback,
		timeout=2000,
		behavior=Button.BEHAVIOR_RELEASE
	)

	options = [
		'Server mode, broadcasts ssid & password to clients',
		'Client mode, receive ssid & password from server',
	]

	# 运行模式三选一
	mode = Utilities.choose_an_option('BLE Config Device Mode', options)
	# mode = MODE_SERVER
	# mode = MODE_CLIENT

	if mode is not None:
		if mode == MODE_SERVER:
			run_bleconfig_server()
		else:
			run_bleconfig_client()


if __name__ == '__main__':
	run_test()
